from copy import deepcopy
import random
import time
from chiplet import QPU
from tool import remove_root_node_from_dag, extract_first_layer
from math import log2, ceil
import networkx as nx
import queue


# ⭐️ 选路
def select_sole_path_for_given_CNOTs(qpu: QPU, cnots: list[list]):
    '''待优化的地方，针对给定的一批CNOT门（不一定是全部的CNOT门），选择唯一的路径\n
    传入一批CNOT门，为这些门请求选择路径\n
    为每个CNOT节点对，计算唯一的路径
    '''
    arr = []
    for cnot in cnots:
        tmp = qpu.get_k_path_with_qubit_id(
            source_qubit=cnot[0], target_qubit=cnot[1], k=1)
        arr.append(tmp[0])
    return arr


# ⭐️ 调度
def schedule_greedy(original_qpu: QPU):
    '''贪心调度，计算最后电路完成时间'''
    qpu: QPU = deepcopy(original_qpu)  # 先拷贝一份
    # ⭐️ 在调度的过程中，量子比特的映射会发生变化
    final_time_dag_greedy = deepcopy(qpu.circuit_dag.dag_nx)

    tmp_dag_greedy = deepcopy(qpu.circuit_dag.dag_nx)
    remove_root_node_from_dag(tmp_dag_greedy)
    has_remaining_cnot = True

    # swap_num = 0
    # ghz_length = 0

    ancilla_num = 0

    total_consume_times = []

    while has_remaining_cnot:
        # 抽取 Front Layer
        cnots_in_window, tmp_dag_with_attrs, node_id_in_window, _has_remaining_cnot = extract_first_layer(
            tmp_dag_greedy)

        # 判断DAG图中是否有剩余的CNOT门没处理
        has_remaining_cnot = _has_remaining_cnot
        if (not has_remaining_cnot):
            break

        # 采样其中一部分CNOT门
        selected_cnots = cnots_in_window
        selected_node_dag_ids = node_id_in_window

        start_time = time.time()

        cnot_paths = select_sole_path_for_given_CNOTs(qpu, selected_cnots)

        # 排序index，调出最短的先满足
        sorted_indices = sorted(range(len(cnot_paths)),
                                key=lambda i: len(cnot_paths[i]))

        # 根据排序后的索引重新排列 selected_cnots 和 selected_node_dag_ids
        sorted_cnots = [selected_cnots[i] for i in sorted_indices]
        sorted_node_dag_ids = [selected_node_dag_ids[i]
                               for i in sorted_indices]
        sorted_paths = [cnot_paths[i] for i in sorted_indices]

        non_conflicting_path_idxs = find_conflict(sorted_paths)

        # print("processed dag：", [sorted_node_dag_ids[i]
        #       for i in non_conflicting_path_idxs])

        for idx in non_conflicting_path_idxs:
            dag_node_id = sorted_node_dag_ids[idx]
            path = sorted_paths[idx]
            max_pre_time = 0
            for pre in final_time_dag_greedy.predecessors(dag_node_id):
                max_pre_time = max(
                    max_pre_time, final_time_dag_greedy.nodes[pre]["process_done_time"])
            # 所用的辅助Qubit的数量
            cross_computing_qubit_num = get_cross_computing_qubit_num(
                path, qpu)
            # move away computing qubits
            final_time_dag_greedy.nodes[dag_node_id]["process_done_time"] = max_pre_time + ceil(log2(
                len(path))) + get_move_away_time(path, qpu)

            # ghz_length += len(path)
            # swap_num += cross_computing_qubit_num

            ancilla_num += len(path)
            # ancilla_num += cross_computing_qubit_num

            diff_idxs = [
                x for x in sorted_indices if x not in non_conflicting_path_idxs]

            for idx1 in diff_idxs:
                dag_node_id1 = sorted_node_dag_ids[idx1]
                path1 = sorted_paths[idx1]
                if (is_two_path_conflict(path, path1)):
                    # 如果他俩冲突的话，就增加一条边，保证DAG的偏序性
                    final_time_dag_greedy.add_edge(
                        dag_node_id, dag_node_id1, weight=1)
        tmp_dag_greedy.remove_nodes_from(
            [sorted_node_dag_ids[i] for i in non_conflicting_path_idxs])

        end_time = time.time()
        total_consume_times.append(end_time - start_time)

    # print("time_greedy: ", sum(total_consume_times))

    max_depth = 0
    # 找出 self.final_depth_dag 出度为0的顶点
    for node in final_time_dag_greedy.nodes():
        if (final_time_dag_greedy.nodes[node]["process_done_time"] > max_depth):
            max_depth = final_time_dag_greedy.nodes[node]["process_done_time"]

    # print("max_time_greedy: ", max_depth)
    return max_depth, ancilla_num


def is_two_path_conflict(path1, path2) -> bool:
    '''判断两条路径是否存在冲突'''
    links_set = set()
    for i in range(len(path1) - 1):
        u, v = path1[i], path1[i + 1]
        link = (min(u, v), max(u, v))
        links_set.add(link)
    for i in range(len(path2) - 1):
        u, v = path2[i], path2[i + 1]
        link = (min(u, v), max(u, v))
        if link in links_set:
            return True
    return False


def find_conflict(paths):
    '''找到多条路径彼此不冲突的，或者冲突路径中最短的那一条'''
    link_count = {}
    path_links = []
    for path in paths:
        links = set((min(path[i], path[i+1]), max(path[i], path[i+1]))
                    for i in range(len(path)-1))
        path_links.append(links)
        for link in links:
            link_count[link] = link_count.get(link, 0) + 1
    conflicts = [any(link_count[link] > 1 for link in links)
                 for links in path_links]
    non_conflicting_indices = [
        i for i, conflict in enumerate(conflicts) if not conflict]
    conflicting_indices = [
        i for i, conflict in enumerate(conflicts) if conflict]
    shortest_conflicting_index = min(
        conflicting_indices, key=lambda i: len(paths[i])) if conflicting_indices else None
    return non_conflicting_indices + ([shortest_conflicting_index] if shortest_conflicting_index is not None else [])


def get_move_away_time(path, qpu: QPU):
    '''获取并行移动qubit所花费的时间'''
    # path - [物理比特, 物理比特, 物理比特...]
    merged_graph = qpu.merged_graph  # 物理比特拓扑
    map_plan = qpu.map_plan  # { 物理比特：计算比特 ... }

    for i in path:
        if (map_plan[i] != None):  # 存在已映射的计算比特，计算移走它花费的时间
            visited = set()
            visited.add(i)
            q = queue.Queue()
            q.put((i, 1))  # 第一个节点，第一层
            while (not q.empty()):
                cur_node = q.get()
                cur_node_id = cur_node[0]
                cur_node_depth = cur_node[1]
                if (map_plan[cur_node_id] == None):  # 找到一个没有被映射的节点
                    return cur_node_depth
                for neighbor in merged_graph.neighbors(cur_node_id):
                    if (neighbor not in visited):
                        visited.add(neighbor)
                        q.put((neighbor, cur_node_depth + 1))
    return 0


def get_cross_computing_qubit_num(path, qpu: QPU):
    '''获取穿过了的量子比特数'''
    count = 0
    for i in path:
        if (qpu.map_plan[i] != None):
            count += 1
    return count
